package com.ckx.config.loadbalancer;

import com.alibaba.cloud.nacos.NacosServiceInstance;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;

@SuppressWarnings("unused")
@Slf4j
public class GrayRobinLoadBalancer extends RoundRobinLoadBalancer {

    private final ObjectProvider<ServiceInstanceListSupplier> instanceListSupplierProvider;

    private final String serviceId;

    public GrayRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> instanceListSupplierProvider,
                                 String serviceId) {
        super(instanceListSupplierProvider, serviceId);
        this.instanceListSupplierProvider = instanceListSupplierProvider;
        this.serviceId = serviceId;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = instanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next()
                .map(serviceInstances -> getInstanceResponse(serviceInstances, request));
    }

    @SuppressWarnings("rawtypes")
    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, Request request) {

        if (instances.isEmpty()) {
            log.warn("No servers available for service: " + serviceId);
            return new EmptyResponse();
        }

        DefaultRequestContext requestContext = (DefaultRequestContext) request.getContext();
        RequestData requestData = (RequestData) requestContext.getClientRequest();
        HttpHeaders headers = requestData.getHeaders();
        String serviceVersion = "version";
        String grayTag = headers.getFirst(serviceVersion);
        if (StringUtils.isBlank(grayTag)) {
            return super.choose(request).block();
        }

        for (ServiceInstance instance : instances) {
            NacosServiceInstance nacosServiceInstance = (NacosServiceInstance) instance;
            Map<String, String> metadata = nacosServiceInstance.getMetadata();
            String value = metadata.getOrDefault(serviceVersion, "");
            if ("gray".equalsIgnoreCase(value)) {
                log.info("服务: {}, 命中灰度节点: {}", serviceId, nacosServiceInstance.getHost()
                        + ":" + nacosServiceInstance.getPort());
                return new DefaultResponse(instance);
            }
        }

        return super.choose(request).block();
    }

}
